"""
数据预处理模块
负责数据加载、清洗和基本转换
"""

import pandas as pd
import numpy as np
import os
import warnings
from sklearn.model_selection import train_test_split

# 忽略警告
warnings.filterwarnings('ignore')


def load_data(train_path, test_path):
    """
    加载训练和测试数据

    参数:
        train_path (str): 训练数据路径
        test_path (str): 测试数据路径

    返回:
        tuple: (train_df, test_df) 训练和测试数据DataFrame
        
    异常:
        FileNotFoundError: 如果数据文件不存在
    """
    # 检查文件是否存在
    if not os.path.exists(train_path):
        raise FileNotFoundError(f"训练数据文件不存在: {train_path}")
    if not os.path.exists(test_path):
        raise FileNotFoundError(f"测试数据文件不存在: {test_path}")
        
    try:
        train_df = pd.read_csv(train_path)
        test_df = pd.read_csv(test_path)

        print(f"训练数据形状: {train_df.shape}")
        print(f"测试数据形状: {test_df.shape}")
        
        # 检查数据是否包含必要的列
        required_columns = ['PassengerId', 'Survived', 'Pclass', 'Name', 'Sex', 'Age', 
                          'SibSp', 'Parch', 'Ticket', 'Fare', 'Cabin', 'Embarked']
        
        missing_train_cols = [col for col in required_columns if col not in train_df.columns]
        missing_test_cols = [col for col in required_columns if col != 'Survived' and col not in test_df.columns]
        
        if missing_train_cols:
            print(f"警告: 训练数据缺少以下列: {missing_train_cols}")
        if missing_test_cols:
            print(f"警告: 测试数据缺少以下列: {missing_test_cols}")

        return train_df, test_df
    except Exception as e:
        print(f"加载数据时出错: {e}")
        raise


def explore_data(df, show_samples=True):
    """
    探索数据基本信息

    参数:
        df (pd.DataFrame): 要探索的数据框
        show_samples (bool): 是否显示数据样本

    返回:
        dict: 包含数据基本信息的字典
    """
    print("=" * 50)
    print("数据探索分析")
    print("=" * 50)
    
    # 基本信息
    info = {
        "shape": df.shape,
        "columns": list(df.columns),
        "dtypes": df.dtypes.to_dict(),
        "missing_values": df.isnull().sum().to_dict(),
        "missing_percentage": (df.isnull().sum() / len(df) * 100).to_dict()
    }

    # 显示基本信息
    print(f"数据形状: {info['shape']}")
    print(f"特征数量: {len(info['columns'])}")
    print(f"样本数量: {info['shape'][0]}")
    
    # 显示数据样本
    if show_samples:
        print("\n数据样本 (前5行):")
        print(df.head())
    
    # 显示数据类型
    print("\n数据类型:")
    for col, dtype in info['dtypes'].items():
        print(f"{col}: {dtype}")
    
    # 显示缺失值统计
    print("\n缺失值统计:")
    has_missing = False
    for col, missing in info['missing_values'].items():
        if missing > 0:
            print(f"{col}: {missing} ({info['missing_percentage'][col]:.2f}%)")
            has_missing = True
    
    if not has_missing:
        print("数据中没有缺失值")
    
    # 显示数值型特征的统计信息
    print("\n数值型特征统计信息:")
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    if len(numeric_cols) > 0:
        print(df[numeric_cols].describe())
    else:
        print("数据中没有数值型特征")
    
    # 显示类别型特征的统计信息
    print("\n类别型特征统计信息:")
    categorical_cols = df.select_dtypes(include=['object']).columns
    if len(categorical_cols) > 0:
        for col in categorical_cols:
            print(f"\n{col} (唯一值数量: {df[col].nunique()}):")
            print(df[col].value_counts().head(10))
    else:
        print("数据中没有类别型特征")
    
    # 如果有Survived列，显示目标变量分布
    if 'Survived' in df.columns:
        print("\n目标变量 (Survived) 分布:")
        survival_counts = df['Survived'].value_counts()
        print(f"幸存: {survival_counts.get(1, 0)} ({survival_counts.get(1, 0) / len(df) * 100:.2f}%)")
        print(f"遇难: {survival_counts.get(0, 0)} ({survival_counts.get(0, 0) / len(df) * 100:.2f}%)")
    
    return info


def handle_missing_values(df):
    """
    处理缺失值

    参数:
        df (pd.DataFrame): 原始数据框

    返回:
        pd.DataFrame: 处理后的数据框
    """
    print("=" * 50)
    print("缺失值处理")
    print("=" * 50)
    
    df_processed = df.copy()
    
    # 记录处理前的缺失值情况
    print("处理前缺失值情况:")
    missing_before = df_processed.isnull().sum()
    for col, missing in missing_before.items():
        if missing > 0:
            print(f"{col}: {missing} ({missing / len(df_processed) * 100:.2f}%)")

    # 1. 年龄(Age) - 按船舱等级(Pclass)和性别(Sex)分组的中位数填充
    if 'Age' in df_processed.columns and df_processed['Age'].isnull().sum() > 0:
        print("\n处理年龄(Age)缺失值...")
        age_by_pclass_sex = df_processed.groupby(['Pclass', 'Sex'])['Age'].median()
        
        for pclass in df_processed['Pclass'].unique():
            for sex in df_processed['Sex'].unique():
                mask = (df_processed['Pclass'] == pclass) & (df_processed['Sex'] == sex) & (df_processed['Age'].isnull())
                if mask.sum() > 0:  # 只有当存在缺失值时才处理
                    try:
                        median_age = age_by_pclass_sex.loc[(pclass, sex)]
                        df_processed.loc[mask, 'Age'] = median_age
                        print(f"  Pclass={pclass}, Sex={sex}: 填充了 {mask.sum()} 个缺失值, 中位数年龄={median_age:.2f}")
                    except KeyError:
                        # 如果某个组合不存在，使用整体中位数
                        global_median = df_processed['Age'].median()
                        df_processed.loc[mask, 'Age'] = global_median
                        print(f"  Pclass={pclass}, Sex={sex}: 填充了 {mask.sum()} 个缺失值, 使用全局中位数={global_median:.2f}")

    # 2. 票价(Fare) - 按船舱等级的中位数填充
    if 'Fare' in df_processed.columns and df_processed['Fare'].isnull().sum() > 0:
        print("\n处理票价(Fare)缺失值...")
        fare_by_pclass = df_processed.groupby('Pclass')['Fare'].median()
        
        for pclass in df_processed['Pclass'].unique():
            mask = (df_processed['Pclass'] == pclass) & (df_processed['Fare'].isnull())
            if mask.sum() > 0:  # 只有当存在缺失值时才处理
                try:
                    median_fare = fare_by_pclass.loc[pclass]
                    df_processed.loc[mask, 'Fare'] = median_fare
                    print(f"  Pclass={pclass}: 填充了 {mask.sum()} 个缺失值, 中位数票价={median_fare:.2f}")
                except KeyError:
                    # 如果某个船舱等级不存在，使用整体中位数
                    global_median = df_processed['Fare'].median()
                    df_processed.loc[mask, 'Fare'] = global_median
                    print(f"  Pclass={pclass}: 填充了 {mask.sum()} 个缺失值, 使用全局中位数={global_median:.2f}")

    # 3. 登船港口(Embarked) - 用最频繁值填充
    if 'Embarked' in df_processed.columns and df_processed['Embarked'].isnull().sum() > 0:
        print("\n处理登船港口(Embarked)缺失值...")
        most_common_embarked = df_processed['Embarked'].mode()[0]
        embarked_missing = df_processed['Embarked'].isnull().sum()
        df_processed['Embarked'].fillna(most_common_embarked, inplace=True)
        print(f"  填充了 {embarked_missing} 个缺失值, 使用最频繁值 '{most_common_embarked}'")

    # 4. 船舱号(Cabin) - 提取首字母，缺失值标记为"Unknown"
    if 'Cabin' in df_processed.columns:
        print("\n处理船舱号(Cabin)缺失值...")
        cabin_missing = df_processed['Cabin'].isnull().sum()
        df_processed['Cabin'] = df_processed['Cabin'].fillna('Unknown')
        df_processed['Cabin_Letter'] = df_processed['Cabin'].str[0]
        print(f"  填充了 {cabin_missing} 个缺失值, 标记为 'Unknown'")
        print(f"  船舱首字母分布: {df_processed['Cabin_Letter'].value_counts().to_dict()}")

    # 记录处理后的缺失值情况
    print("\n处理后缺失值情况:")
    missing_after = df_processed.isnull().sum()
    still_missing = False
    for col, missing in missing_after.items():
        if missing > 0:
            print(f"{col}: {missing} ({missing / len(df_processed) * 100:.2f}%)")
            still_missing = True
    
    if not still_missing:
        print("所有缺失值已处理完成")
    else:
        print("警告: 仍有部分缺失值未处理")
        
    return df_processed


def split_data(df, target_col='Survived', test_size=0.2, random_state=42):
    """
    将数据分割为训练集和验证集

    参数:
        df (pd.DataFrame): 完整的数据集
        target_col (str): 目标列名
        test_size (float): 测试集比例
        random_state (int): 随机种子

    返回:
        tuple: (X_train, X_val, y_train, y_val)
        
    异常:
        ValueError: 如果目标列不在数据框中
    """
    print("=" * 50)
    print("数据集划分")
    print("=" * 50)
    
    # 检查目标列是否存在
    if target_col not in df.columns:
        raise ValueError(f"目标列 '{target_col}' 不在数据框中")
    
    try:
        X = df.drop(target_col, axis=1)
        y = df[target_col]

        X_train, X_val, y_train, y_val = train_test_split(
            X, y, test_size=test_size, random_state=random_state, stratify=y
        )

        print(f"训练集形状: {X_train.shape}")
        print(f"验证集形状: {X_val.shape}")
        
        # 检查划分后的数据分布
        print("\n训练集目标变量分布:")
        train_survived = y_train.value_counts(normalize=True)
        print(f"幸存: {train_survived.get(1, 0) * 100:.2f}%")
        print(f"遇难: {train_survived.get(0, 0) * 100:.2f}%")
        
        print("\n验证集目标变量分布:")
        val_survived = y_val.value_counts(normalize=True)
        print(f"幸存: {val_survived.get(1, 0) * 100:.2f}%")
        print(f"遇难: {val_survived.get(0, 0) * 100:.2f}%")

        return X_train, X_val, y_train, y_val
    except Exception as e:
        print(f"分割数据时出错: {e}")
        raise
